package main

import (
    "log"
    "fmt"
    "time"
    "crypto/sha1"
    "crypto/rsa"
    zmq "github.com/pebbe/zmq4"
    "github.com/vmihailenco/msgpack"
)

func (m *Minion) startPubChannel() {

    var zmq_filtering bool = false

    /* Step 1. Create connection */
    client, err := zmq.NewSocket(zmq.SUB)
    if err != nil {
	panic(err)
    }
    var pubAddrStr = fmt.Sprintf("tcp://%s:%d", m.Addr, m.PubPort)
    log.Printf("Connect publisher on %s", pubAddrStr)
    client.Connect(pubAddrStr)
    client.SetLinger(-1)
    if zmq_filtering {
        //self._socket.setsockopt(zmq.SUBSCRIBE, b"broadcast")
        //self._socket.setsockopt(zmq.SUBSCRIBE, salt.utils.stringutils.to_bytes(self.hexid))
    } else {
        client.SetSubscribe("")
    }
    // if _id {
    //    self._socket.setsockopt(zmq.IDENTITY, salt.utils.stringutils.to_bytes(_id))
    //}

    /* Setup Keep-Alive options */
    //if hasattr(zmq, "TCP_KEEPALIVE"):
    //        self._socket.setsockopt(zmq.TCP_KEEPALIVE, tcp_keepalive)
    //        self._socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, tcp_keepalive_idle)
    //        self._socket.setsockopt(zmq.TCP_KEEPALIVE_CNT, tcp_keepalive_cnt)
    //        self._socket.setsockopt(zmq.TCP_KEEPALIVE_INTVL, tcp_keepalive_intvl)

    /* Step 2. Set re-connect duration, sec. */
    var recon_default  = 600 // 5 min.
    var recon_max = 3600 // 1 hour
    var recon_delay = randint(recon_default, recon_default + recon_max)
    log.Printf("Generated random reconnect delay between '%dms' and '%dms' (%d)", recon_delay, recon_delay + recon_max, recon_delay)
    log.Printf("Setting zmq_reconnect_ivl to '%dms'", recon_delay)
    client.SetReconnectIvl(time.Second * time.Duration(recon_delay))
    log.Printf("Setting zmq_reconnect_ivl_max to '%dms'", recon_delay + recon_max)
    client.SetReconnectIvlMax(time.Second * time.Duration(recon_max))

    for {
        msg, _ := client.RecvMessage(0)
        var raw []byte = []byte(msg[0])

        var resp map[string]interface{} = make(map[string]interface{})
        _ = msgpack.Unmarshal(raw, &resp)
        log.Printf("pub - %+v", resp)

        var txt string = resp["load"].(string)
        raw2, _ := m.c.Decode(txt)
        log.Printf("load = %s", raw2)

        //
        var raw3 []byte = []byte(raw2)
        var resp2 map[string]interface{} = make(map[string]interface{})
        err = msgpack.Unmarshal(raw3, &resp2)
        log.Printf("resp = %v %v", err, resp2)

        //
        c := NewCommand()
        c.SetFun(resp2["fun"].(string))
        c.SetJid(resp2["jid"].(string))

        //
        e := NewEngine()
        e.execute(*m, c)

    }

}

type Minion struct {
    Id          string         /* Minion ID                */
    client     *zmq.Socket     /* ZeroMQ connection        */
    c          *Crypticle      /* Crypticle service        */
    Addr        string         /* Salt master address      */
    ReqPort     uint16         /* Salt master req port     */
    PubPort     uint16         /* Salt master pub port     */
}

func NewMinion() Minion {
    return Minion{
        Addr: "127.0.0.1",
        ReqPort: 4505,
        PubPort: 4506,
    }
}

func (m Minion) SetID(id string) {
    m.Id = id
}

func (m Minion) answer(jid string, r interface{}) {
    var payload map[string]interface{} = make(map[string]interface{})
    var args []string = []string{}
    payload["cmd"] = "_return"
    payload["id"] = "spy"
    payload["success"] = true
    payload["return"] = r
    payload["retcode"] = 0
    payload["jid"] = jid
    payload["fun"] = "test.ping"
    payload["fun_args"] = args
    payload["user"] = "root"
    payload["_stamp"] = "AAAA"
    resp, err := m.executeCommand(payload, true)
    log.Printf("answer = %+v err = %s", resp, err)
}

func (m *Minion) connect() error {
    var connStr string = fmt.Sprintf("tcp://%s:%d", m.Addr, m.ReqPort)
    var err error
    log.Printf("Create connection on %s", connStr)
    m.client, err = zmq.NewSocket(zmq.REQ)
    if err != nil {
	return err
    }
    m.client.Connect(connStr)
    return nil
}

func (m *Minion) executeCommand(payload map[string]interface{}, secure bool) (map[string]interface{}, error) {
    var err error
    var raw1 []byte // Input RAW packet
    var raw2 []byte // Output RAW packet

    /* Step 1. Prepare message */
    if secure {
        log.Printf("ReqChannel send crypt load=%+v", payload)
        var req []byte // RAW packet
        /* Step 1. Crypt AES payload */
        req, err = msgpack.Marshal(&payload)
        if err != nil {
            return nil, err
        }
        pkt, err := m.c.Encode(PICKLE_PAD + string(req))
        if err != nil {
            return nil, err
        }
        /* Step 2. Wrap packet */
        var wrap map[string]interface{} = make(map[string]interface{})
        wrap["enc"] = "aes"
        wrap["load"] = pkt
        raw1, err = msgpack.Marshal(&wrap)
        if err != nil {
            return nil, err
        }
    } else {
        log.Printf("ReqChannel send clear load=%+v", payload)
        /* Step 1. Wrap packet */
        var wrap map[string]interface{} = make(map[string]interface{})
        wrap["enc"] = "clear"
        wrap["load"] = payload
        raw1, err = msgpack.Marshal(&wrap)
        if err != nil {
            panic(err)
            return nil, err
        }
    }

    /* Step 2. Network I/O operation */
    var msg []string
    m.client.SendMessage(raw1)
    msg, err = m.client.RecvMessage(0)
    if err != nil {
        panic(err)
        return nil, err
    }

    /* Step 3. Parse response */
//    var resp map[string]interface{} = make(map[string]interface{})
    var resp interface{}
    if secure {
        raw2 = []byte(msg[0])
        log.Printf("secure: raw2 = %s", raw2)
        pkt, err := m.c.Decode(string(raw2))
        if err != nil {
            return nil, err
        }
        err = msgpack.Unmarshal([]byte(pkt), &resp)
        if err != nil {
            return nil, err
        }
        log.Printf("resp = %+v", err, resp)
    } else {
        raw2 = []byte(msg[0])
        log.Printf("clear: raw2 = %s", raw2)
        err = msgpack.Unmarshal(raw2, &resp)
        if err != nil {
            panic(err)
            return nil, err
        }
        log.Printf("resp = %+v", err, resp)
    }

    /* Return result */
    if err_ret, ok := resp.(string); ok {
        panic(err_ret)
    } else if map_ret, ok := resp.(map[string]interface{}); ok {
        return map_ret, err
    } else {
        log.Printf("resp = %+v", resp)
        panic("some error")
    }

}

func (m *Minion) start() error {
    var err error

    /* Step 1. Create connection */
    m.connect()

    m_rsa := restoreKey()

    var id string = "spy"

    var payload map[string]interface{} = make(map[string]interface{})
    payload["cmd"] = "_auth"
    payload["id"] = id
    payload["nonce"] = ""
    payload["pub"] = m_rsa.pubKey

    /* Step 1. Authorize routine */
    var resp map[string]interface{}
    var publish_port uint16 = 0
    for {
        /* Step 1. Ask authroze request */
        resp, err = m.executeCommand(payload, false)
        if err != nil {
            panic(err)
        }
        /* Step 2. Check response */
        if value, ok := resp["publish_port"].(uint16); ok {
            log.Printf("Auth complete")
            publish_port = value
            break
        }
        /* Step 3. Wait */
        log.Printf("Waiting 10 seconds before retry.")
        time.Sleep(10 * time.Second)
    }

    /* Temp - move at publisher */
    var secretKey []byte = []byte(resp["aes"].(string))
    //log.Printf("aes = %+v", secretKey)

    aesKey, err := rsa.DecryptOAEP(sha1.New(), nil, m_rsa.key, secretKey, nil)
    if err != nil {
	panic(err)
    }
    log.Printf("use AES key = %s",  aesKey)

    m.c = NewCrypticle(string(aesKey))

    /* Step 3. Subscribe on publisher */
    log.Printf("Start publisher subscribe %d", publish_port)
    m.startPubChannel()

    /* Step 2. Operating */
    for {

        time.Sleep(10 * time.Second)
    }

    return err
}
